from machine import Pin, SoftI2C
import time
import math

# 比例增益，用于加速度计/磁强计的收敛
Kp=0.8
# 积分增益，用于陀螺仪偏见的衔接
Ki=0.001
# 采样周期的一半
halfT=0.004
# 四元数元素，代表估计方向
q0=1
q1=0
q2=0
q3=0
# 按比例缩小的积分误差
exInt=0
eyInt=0
ezInt=0

def IMUupdate(gx, gy, gz, ax, ay, az):
    """更新IMU数据并计算姿态角    
    参数:
        gx, gy, gz: 陀螺仪x,y,z轴角速度
        ax, ay, az: 加速度计x,y,z轴加速度        
    返回:
        包含pitch, roll, 加速度和角速度的数组
    """
    K = 0.5
    a = [0, 0, 0, 0, 0, 0, 0, 0]
    global Kp, Ki, halfT, q0, q1, q2, q3, exInt, eyInt, ezInt
    # 单位化加速度数据
    norm = math.sqrt(ax * ax + ay * ay + az * az)
    if norm != 0:
        ax = ax / norm
        ay = ay / norm
        az = az / norm
    # 估计方向的重力
    vx = 2 * (q1 * q3 - q0 * q2)
    vy = 2 * (q0 * q1 + q2 * q3)
    vz = q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3
    # 计算误差
    ex = (ay * vz - az * vy)
    ey = (az * vx - ax * vz)
    ez = (ax * vy - ay * vx)
    # 积分误差
    exInt = exInt + ex * Ki
    eyInt = eyInt + ey * Ki
    ezInt = ezInt + ez * Ki
    # 调整后的陀螺仪测量
    gx = gx + Kp * ex + exInt
    gy = gy + Kp * ey + eyInt
    gz = gz + Kp * ez + ezInt
    # 更新四元数
    q0 = q0 + (-q1 * gx - q2 * gy - q3 * gz) * halfT
    q1 = q1 + (q0 * gx + q2 * gz - q3 * gy) * halfT
    q2 = q2 + (q0 * gy - q1 * gz + q3 * gx) * halfT
    q3 = q3 + (q0 * gz + q1 * gy - q2 * gx) * halfT
    # 归一化四元数
    norm = math.sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3)
    if norm != 0:
        q0 = q0 / norm
        q1 = q1 / norm
        q2 = q2 / norm
        q3 = q3 / norm
    
    # 计算 Pitch 和 Roll
    # 第一次计算（基于四元数）
    Pitch = math.asin(-2 * q1 * q3 + 2 * q0 * q2) * 57.3  # pitch，转换为度数
    Roll = 0  # 初始化 Roll
    denom = -2 * q1 * q1 - 2 * q2 * q2 + 1
    if denom != 0:
        Roll = math.atan((2 * q2 * q3 + 2 * q0 * q1) / denom) * 57.3  # roll，转换为度数
   
    # 第二次计算（基于加速度计）
    if ay * ay + az * az != 0:
        a[2] = -math.atan(ax / math.sqrt(ay * ay + az * az)) * 57.2957795
    if ax * ax + az * az != 0:
        a[3] = math.atan(ay / math.sqrt(ax * ax + az * az)) * 57.2957795
    a[4] = gx
    a[5] = gy
    a[6] = gz
    # 最终融合（互补滤波）
    a[0] = -K * Pitch - (1 - K) * a[2]  # 融合后的Pitch
    a[1] = K * Roll + (1 - K) * a[3]    # 融合后的Roll
    return a


class accel():
    gyro_zero_offset = [0, 0, 0]
    def __init__(self, i2c, addr=0x68):
        self.iic = i2c
        self.addr = addr
        self.acc_bias = [0, 0, 0]
        # 寄存器常量定义（新增）
        PWR_MGMT_1 = 0x6B
        CONFIG = 0x1A
        GYRO_CONFIG = 0x1B
        ACCEL_CONFIG = 0x1C
        ACCEL_CONFIG2 = 0x1D
    
        # 修改后的初始化配置
        self.iic.writeto_mem(self.addr, PWR_MGMT_1, b'\x00')  # 唤醒
        self.iic.writeto_mem(self.addr, CONFIG, b'\x03')      # 陀螺仪DLPF_CFG=3
        self.iic.writeto_mem(self.addr, ACCEL_CONFIG2, b'\x03') # 加速度计DLPF_CFG=3
        self.iic.writeto_mem(self.addr, ACCEL_CONFIG, b'\x08')  # ±4g
        self.iic.writeto_mem(self.addr, GYRO_CONFIG, b'\x08')    # ±500dps
    
    def _read_word(self, addr):
        high = self.iic.readfrom_mem(self.addr, addr, 1)[0]
        low = self.iic.readfrom_mem(self.addr, addr+1, 1)[0]
        val = (high << 8) + low
        return val - (0x10000 if val > 0x7FFF else 0)
    
    def gyro_zero_offset_gy(self, samples=100):
        # 同时校准加速度计和陀螺仪
        ax_sum = ay_sum = az_sum = 0
        gx_sum = gy_sum = gz_sum = 0
        
        for _ in range(samples):
            ax_sum += self._read_word(0x3B)
            ay_sum += self._read_word(0x3B+2)
            az_sum += self._read_word(0x3B+4)
            gx_sum += self._read_word(0x43)
            gy_sum += self._read_word(0x45)
            gz_sum += self._read_word(0x47)
            time.sleep_ms(10)
            
        self.acc_bias = [
            ax_sum/samples,
            ay_sum/samples,
            az_sum/samples - 8192  # Z轴补偿1g
        ]
        self.gyro_zero_offset = [
            gx_sum/samples,
            gy_sum/samples,
            gz_sum/samples
        ]
    def get_values(self):
        # 读取校准后的原始数据
        ax = (self._read_word(0x3B) - self.acc_bias[0]) / 8192.0
        ay = (self._read_word(0x3B+2) - self.acc_bias[1]) / 8192.0
        az = (self._read_word(0x3B+4) - self.acc_bias[2]) / 8192.0
        gx = (self._read_word(0x43) - self.gyro_zero_offset[0]) / 65.5
        gy = (self._read_word(0x45) - self.gyro_zero_offset[1]) / 65.5
        gz = (self._read_word(0x47) - self.gyro_zero_offset[2]) / 65.5
        # 直接读取原始值（硬件已滤波）
        ax = (self._read_word(0x3B) - self.acc_bias[0]) / 8192.0
        ay = (self._read_word(0x3B+2) - self.acc_bias[1]) / 8192.0
        az = (self._read_word(0x3B+4) - self.acc_bias[2]) / 8192.0
        gx = (self._read_word(0x43) - self.gyro_zero_offset[0]) / 65.5
        gy = (self._read_word(0x45) - self.gyro_zero_offset[1]) / 65.5
        gz = (self._read_word(0x47) - self.gyro_zero_offset[2]) / 65.5        
        tmp = self._read_word(0x41) / 340.0 + 36.53  # 转换为摄氏度
        return {
            'AcX': ax,  # 直接使用原始值
            'AcY': ay,
            'AcZ': az,
            'GyX': gx,
            'GyY': gy,
            'GyZ': gz,
            'Tmp': tmp
        }
from setcfg import DAT
print("初始化IMU6050陀螺仪 scl=%d,sda=%d"%(DAT["mpu6050"]["scl"],DAT["mpu6050"]["sda"]))

try:  
  i2c = SoftI2C(scl=Pin(DAT["mpu6050"]["scl"]), sda=Pin(DAT["mpu6050"]["sda"]))
  acc = accel(i2c)  # 正确类名
  acc.gyro_zero_offset_gy(100)  # 正确校准方法
except Exception as e:
  print("陀螺仪MPU 6050异常:" +str(e))
